package router

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"net"
	"net/http"
	"os/exec"
	"regexp"
	"strings"
)

type ExecHandler struct {
}

func NewExecHandler() *ExecHandler {
	return &ExecHandler{}
}

func (handler *ExecHandler) RegisterRoutes(r *gin.RouterGroup) {
	r.GET("/ping", handler.Ping)
	r.GET("/exec_sh", handler.ExecSh)
	r.GET("/exec_cmd", handler.ExecCmd)
}

func (handler *ExecHandler) Ping(ctx *gin.Context) {
	ip := ctx.Query("ip")

	if !validIP(ip) {
		ctx.String(http.StatusBadRequest, "Invalid IP address")
		return
	}

	// 创建一个 Command 结构体，指定要执行的命令及其参数
	cmd := exec.Command("bash", "-c", "ping -c 1 " + ip)	//mac
	//cmd := exec.Command("cmd","/c", "ping", "/n", "1", ip)	//windows

	// 使用 CombinedOutput 方法执行命令并捕获标准输出和标准错误
	output, err := cmd.CombinedOutput()
	if err != nil {
		ctx.String(http.StatusOK, "命令执行失败:", err)
	}

	// 将输出回显给前端
	ctx.String(http.StatusOK, string(output))
}

func (handler *ExecHandler) ExecSh(ctx *gin.Context) {
	shellName := ctx.Query("shellName")
	if !isShellNameSafe(shellName) {
		ctx.String(http.StatusBadRequest, "Invalid shellName")
		return
	}

	cmd := exec.Command("bash", "-c",
		fmt.Sprintf("cd %s && /bin/bash %s",
		"/Users/dxm/workspace/go-sec-code/uploads", shellName))

	// 使用 CombinedOutput 方法执行命令并捕获标准输出和标准错误
	output, err := cmd.CombinedOutput()
	if err != nil {
		ctx.String(http.StatusBadRequest, "命令执行失败:", err)
	}

	// 将输出回显给前端
	ctx.String(http.StatusOK, string(output))
}

func (handler *ExecHandler) ExecCmd(ctx *gin.Context) {
	cmd := ctx.Query("cmd")
	if !isExecCmdLegal(cmd) || !isExecArgLegal(cmd) {
		ctx.String(http.StatusBadRequest, "Invalid cmd")
		return
	}

	command := exec.Command("bash", "-c", cmd)

	// 使用 CombinedOutput 方法执行命令并捕获标准输出和标准错误
	output, err := command.CombinedOutput()
	if err != nil {
		ctx.String(http.StatusBadRequest, "命令执行失败:", err)
	}

	// 将输出回显给前端
	ctx.String(http.StatusOK, string(output))
}

// 验证 IP 地址
func validIP(ip string) bool {
	return net.ParseIP(ip) != nil
}

func isShellNameSafe(filename string) bool {
	// 只允许字母、数字和'.'
	matched, _ := regexp.MatchString(`^[a-zA-Z0-9.]+$`, filename)
	return matched
}

// isExecCmdLegal 检查命令是否在白名单中
func isExecCmdLegal(cmd string) bool {
	whitelistCmd := []string{"ping", "curl"} // 视业务情况，自行设置命令白名单
	for _, safeCmd := range whitelistCmd {
		if strings.HasPrefix(cmd, safeCmd + " ") {
			return true
		}
	}
	return false
}

// isExecArgLegal 检查执行参数是否合法，只允许空格、数字、大小写字母
func isExecArgLegal(text string) bool {
	// 使用正则表达式定义白名单字符集
	regex := `^[\s0-9a-zA-Z\-]*$`
	matched, _ := regexp.MatchString(regex, text)
	return matched
}